package org.zlb.algorithm.tree;

/**
 * 字典树
 * 
 * @author zhoulingbo
 * @date 2021/09/18
 */
public class Trie {

    private char ch; // 仅支持英文a-z
    private int[] counts = new int[26];
    private Trie[] chidren = new Trie[26];
    
    public Trie() {
        
    }
    
    public Trie(char ch) {
        this.ch = ch;
    }
    
    public void insert(String key) {
        char[] arr = key.toCharArray();
        insert(arr, 0);
    }
    
    void insert(char[] arr, int i) {
        if (i >= arr.length) {
            return;
        }
        
        int index = arr[i] - 'a';
        this.counts[index]++;
        Trie child = chidren[index];
        if (child == null) {
            child = new Trie(arr[i]);
            this.chidren[index] = child;
        }
        
        int j = i + 1;
        child.insert(arr, j);
    }
    
    public boolean exist(String prefix) {
        char[] arr = prefix.toCharArray();
        return exist(arr, 0);
    }
    
    boolean exist(char[] arr, int i) {
        int index = arr[i] - 'a';
        if (this.counts[index] < 1)
            return false;
        
        if (i == arr.length - 1)
            return true;
        
        Trie child = chidren[index];
        int j = i + 1;
        return child.exist(arr, j);
    }
    
    public int count(String prefix) {
        char[] arr = prefix.toCharArray();
        return count(arr, 0);
    }
    
    int count(char[] arr, int i) {
        int index = arr[i] - 'a';
        if (this.counts[index] < 1)
            return 0;
        
        if (i == arr.length - 1)
            return this.counts[index];
        
        Trie child = chidren[index];
        int j = i + 1;
        return child.count(arr, j);
    }

    public char getCh() {
        return ch;
    }

    public void setCh(char ch) {
        this.ch = ch;
    }

    public int[] getCounts() {
        return counts;
    }

    public void setCounts(int[] counts) {
        this.counts = counts;
    }

    public Trie[] getChidren() {
        return chidren;
    }

    public void setChidren(Trie[] chidren) {
        this.chidren = chidren;
    }
    
}
